import matplotlib.gridspec
import numpy as np 
import cv2, matplotlib
from pathlib import Path
import imageio.v3 as iio 
import matplotlib.pyplot as plt
from tkinter.filedialog import askopenfilenames
from scipy.optimize import curve_fit

def gaussian(x, *params):
    A, mu, sigma, offset = params
    gaussian = np.exp(-(x-mu)**2/(2*sigma**2)) / (sigma * (2*np.pi)**0.5)
    return A * gaussian + offset

rms = lambda X: np.sqrt(np.mean(np.square(X)))


def main(path, gray_value_thresh=2000, fwhm_thresh=50, pixel_size=3.25, n_lines=20):

    img = iio.imread(str(path)).squeeze().astype('uint16')
    img = cv2.medianBlur(img, 3)
    img[img > gray_value_thresh] = rms(img)

    fig, ax = plt.subplots(2, 1, figsize=[5,4], height_ratios=[1,3])
    fig.subplots_adjust(hspace=0.05)
    fig.suptitle(f'Line Width Analysis\n({path.stem})')
    ax[0].axis('off')
    ax[0].imshow(img.T)

    v, h = img.shape
    xs = np.arange(h) * pixel_size
    
    widths = []
    for i in range(int(v/n_lines)):
        img_i = img[i*n_lines:(i+1)*n_lines]
        ys = np.sum(img_i, axis=0)

        try:
            popt, _ = curve_fit(gaussian, xs, ys,
                                    p0=[np.sum(ys)*pixel_size, np.mean(xs), np.ptp(xs)/10, 0],
                                    bounds=([0, -np.inf, 0, -np.inf],
                                            [np.inf, np.inf, np.inf, np.inf]),
                                    )
            fwhm = popt[2] * 2 * (2 * np.log(2))**0.5
            if fwhm > fwhm_thresh:
                fwhm = 0
            widths.append(fwhm)
            print(f'FWHM_{i}: {fwhm:.2f} um')


        except Exception as e:
            print(e)
            widths.append(0)

    widths=np.array(widths)
    widths_eff = widths[widths>0]
    widths[widths==0] = rms(widths_eff)

    ax[1].plot(widths)
    ax[1].set(xlabel='Section', ylabel='FWHM ($\mu m$)')
    ax[1].set_title(f'FWHM (RMS) = {rms(widths):.2f} um  ', y=0.8, loc='right')
    fig.savefig(str(path.with_suffix('.png')))
    plt.show()


if __name__ == '__main__':
    files = askopenfilenames(filetypes=[('tif', ['tif', 'tiff'])])
    if len(files) > 1:
        matplotlib.use('agg')
    for file in files:
        try:
            main(Path(file))
        except Exception as e:
            print(e)